package in.ac.iitb.cse.nlp.postagger.data;

import java.util.HashMap;
import java.util.Map;

public class TransitionMatrix implements Matrix {

	// Internal Map which stores the transition table
	protected Map<Tag, HashMap<Tag, Integer>> transitionTable = new HashMap<Tag, HashMap<Tag, Integer>>();

	public TransitionMatrix() {

	}

	/**
	 * Updates/Increments Count of Tags
	 **/
	public Tag updateCount(String stag) {
		Tag tag = DataMaps.tags.get(stag);
		if (tag == null) {
			tag = new Tag(stag);
			DataMaps.tags.put(stag, tag);
		} else {
			tag.incrementCount();
		}
		return tag;
	}

	/**
	 * Updates/Increments N-gram Count of Tag Updates/Increments
	 * Count(Tn+1=tag2/Tn-1..Tn-k=tag1) currentTag = SINGLE TAG ONLY prevTags =
	 * COMPOSITE PREVIOUS K-1 TAGS FOR K GRAM
	 **/
	// Notion used: tag1 is followed by tag2
	public void update(String prevTags, String currentTag) {
		Tag tagsPrev = DataMaps.tags.get(prevTags);
		if (tagsPrev == null) {
			System.err.println("Howdy !!! How can forget to add PrevTags "
					+ prevTags + " !!!");
			tagsPrev = updateCount(prevTags);
		}
		HashMap<Tag, Integer> map = transitionTable.get(tagsPrev);
		if (map == null) {
			map = new HashMap<Tag, Integer>();
			transitionTable.put(tagsPrev, map);
		}

		Tag tagCurrent = DataMaps.tags.get(currentTag);
		/* If this tag not encountered, create & add it to the first level HMap */
		if (tagCurrent == null) {
			tagCurrent = new Tag(currentTag);
			DataMaps.tags.put(currentTag, tagCurrent);
		}
		/* For sure tag1 cannot be null, update the count */
		Integer count = map.get(tagCurrent);
		if (count == null) {
			map.put(tagCurrent, 1);
		} else {
			map.put(tagCurrent, count + 1);
		}
	}

	/**
	 * 
	 * Returns Count n(Tn+1=currentTag/Tn-1..Tn-k=prevTags) currentTag = SINGLE
	 * TAG ONLY prevTags = COMPOSITE PREVIOUS K-1 TAGS FOR K GRAM
	 **/
	public int getCount(String prevTags, String currentTag) {
		Tag tagsPrev = DataMaps.tags.get(prevTags);
		if (tagsPrev != null) {
			HashMap<Tag, Integer> hashMap = transitionTable.get(tagsPrev);
			if (hashMap != null) {
				Tag tagCurrent = DataMaps.tags.get(currentTag);
				if (tagCurrent != null) {
					Integer integer = hashMap.get(tagCurrent);
					if (integer != null) {
						return integer;
					}
				}
			}
		}
		return 0;
	}

	/**
	 * 
	 * Returns Prob P(Tn+1=currentTag/Tn-1..Tn-k=prevTags) currentTag = SINGLE
	 * TAG ONLY prevTags = COMPOSITE PREVIOUS K-1 TAGS FOR K GRAM
	 **/
	public double getProbability(String prevTags, String currentTag) {
		int countOfColPrecededByRow = getCount(prevTags, currentTag);
		if (countOfColPrecededByRow != 0) {
			Tag tagPrev = DataMaps.tags.get(prevTags);
			if (tagPrev == null) {
				System.err.println("Howdy !!! It is miracle");
				return 0;
			}
			return ((double) countOfColPrecededByRow / (double) tagPrev
					.getCount());
		}
		return 0;
	}

	public int getProbabilityEx(String prevTags, HashMap<Tag, Integer> hashMap) {
		Tag tagsPrev = DataMaps.tags.get(prevTags);
		int countTotal = 0;
		if (tagsPrev != null) {
			hashMap = transitionTable.get(tagsPrev);
			countTotal = tagsPrev.getCount();
		}
		return countTotal;
	}

	public double getProbability_lapsmoothing(String prevTags, String currentTag) {
		int countOfColPrecededByRow = getCount(prevTags, currentTag);
		countOfColPrecededByRow += 1; // Laplace Smoothing
		if (countOfColPrecededByRow != 0) {
			int countTotal = 0;
			Tag tagPrev = DataMaps.tags.get(prevTags);
			if (tagPrev != null)
				countTotal = tagPrev.getCount();
			countTotal += Config.nTotalTags; // Laplace Smoothing
			return ((double) countOfColPrecededByRow / (double) tagPrev
					.getCount());
		}
		return 0;
	}

	@Override
	public String toString() {
		StringBuffer buffer = new StringBuffer();
		for (Tag Prevtags : transitionTable.keySet()) {
			for (Tag currentTag : transitionTable.get(Prevtags).keySet()) {
				buffer.append("("
						+ Prevtags
						+ ","
						+ currentTag
						+ ")="
						+ transitionTable.get(Prevtags).get(currentTag)
						+ ":"
						+ getProbability(Prevtags.getString(),
								currentTag.getString()));
			}
			buffer.append("\n");
		}
		return buffer.toString();
	}
}